"""
    create by IAmFiveHigh on 2024/8/16
"""


class Trie:
    # 字典树 || 前缀树
    class Node:
        def __init__(self):
            self.is_word = False
            self.map = dict()

    def __init__(self):
        self._root = Trie.Node()
        self._size = 0

    def get_size(self) -> int:
        return self._size

    def add(self, element):
        cur = self._root
        for char in element:
            #  如果当前字符在当前结点的字典里不存在
            if char not in cur.map.keys():
                # 创建新的结点
                cur.map[char] = Trie.Node()
            cur = cur.map[char]
        # 如果循环结束结尾is_word为False 说明这是一个新词
        if cur.is_word is False:
            cur.is_word = True
            self._size += 1

    def contains(self, element):
        cur = self._root
        for char in element:
            if char not in cur.map.keys():
                # 如果字典中没有下一个char，说明当前树中没有当前词
                return False
            cur = cur.map[char]
        # 循环结束后，查看最后词是不是is_word结尾
        return cur.is_word

    # 查找是否有前缀为prefix的词
    def is_prefix(self, prefix):
        cur = self._root
        for char in prefix:
            if char not in cur.map.keys():
                return False
            cur = cur.map[char]
        return True


if __name__ == '__main__':
    trie1 = Trie()
    trie1.add("pen")
    trie1.add("pen")
    if trie1.contains("pen") and not trie1.contains("penapple"):
        print("字典树添加和查找没问题")
    else:
        print("字典树添加和查找有bug")

    if trie1.is_prefix("pen") and not trie1.is_prefix("pan"):
        print("字典树is_prefix方法没问题")
    else:
        print("字典树is_prefix方法有bug")
